from pybrain.structure.modules.svmunit import SVMUnit
from pybrain.datasets import SupervisedDataSet
from pybrain.supervised.trainers.svmtrainer import SVMTrainer
from model import *

class SVM():
    inputs = 0
    outputs = 0
    net = None
    ds = None
    trainer = None

    def __init__(self, photos, inputs=0, outputs=0):
        self.inputs = inputs
        self.outputs = outputs
        self.net = SVMUnit(inputs, outputs)
        self.ds = SupervisedDataSet(inputs, outputs)
        for p in photos:
            psds = p.genSupervisedDataset()
            self.ds.addSample(psds[0], psds[1])
        self.trainer = SVMTrainer(svmunit=self.net, dataset=self.ds, plot=True)

    def train(self):
        self.trainer.train()

    def predict(self,p):
        tempds = SupervisedDataSet(self.inputs, self.outputs)
        tempds.addSample(p.genSupervisedDataset()[0], (0,))
        return self.net.activateOnDataset(tempds)

    def load(self, netfile):
        self.net.loadModel(netfile.getFile())

    def save(self, filename):
        self.net.saveModel(filename)
        #s = SVMNets(inputs=self.inputs, outputs=self.outputs, filename=filename)
        #flush()
        #self.net.saveModel(s.getFile())
        return s
        

        